"""
Configuration file for DAT-SNet model
"""
import os
import torch
from datetime import datetime

# Path configurations
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(PROJECT_ROOT, 'data')
SLEEP_EDF_DIR = os.path.join(DATA_DIR, '/home/Wsh/ZYT/Sleep-EDF-20/')
CHECKPOINT_DIR = os.path.join(PROJECT_ROOT, 'checkpoints')
LOGS_DIR = os.path.join(PROJECT_ROOT, 'logs')
RESULTS_DIR = os.path.join(PROJECT_ROOT, 'results')

# Create directories if they don't exist
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
os.makedirs(LOGS_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)

# Dataset configurations
NUM_CLASSES = 5  # W, N1, N2, N3, REM
INPUT_LENGTH = 3000  # 30s at 100Hz
SAMPLING_RATE = 100  # 100Hz


# Model configurations
class ModelConfig:
    # General
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Dataset parameters
    sequence_length = 10  # Number of consecutive epochs to consider for contextual information

    # Temporal Branch
    # 更新卷积核大小，专为睡眠EEG节律优化
    temporal_kernel_sizes = [5, 15, 30, 75, 150]  # 优化后的核大小
    temporal_filters = 16

    # 新增加的增强特性配置
    use_residual = True  # 是否使用残差连接
    use_attention = True  # 是否使用自注意力机制
    use_sble = True  # 是否使用SBLE特征提取
    use_eemd = True  # 是否使用EEMD特征提取

    # Spectral Branch
    wavelet_family = 'db4'  # Daubechies 4
    wavelet_level = 5
    spectral_filters = 64

    # Channel attention
    se_reduction = 16  # Reduction ratio in SE-Block

    # Temporal attention
    bidirectional = True
    hidden_size = 128
    num_layers = 2
    attention_heads = 8
    dropout = 0.5


# Training configurations
class TrainingConfig:
    # Basic settings
    batch_size = 128
    num_epochs = 50
    learning_rate = 1e-3
    weight_decay = 1e-5

    # Optimizer settings
    optimizer = 'adamw'

    # Learning rate scheduler
    lr_scheduler = 'one_cycle'
    lr_scheduler_patience = 5
    lr_scheduler_factor = 0.5

    # Early stopping
    early_stopping_patience = 10

    # Class weights for handling imbalance (W, N1, N2, N3, REM)
    class_weights = [1.0, 2.5, 1.0, 1.0, 1.5]

    # Component-wise training
    feature_extraction_epochs = 100
    sequence_learning_epochs = 100

    # Cross-validation
    n_folds = 20  # 20-fold cross validation for Sleep-EDF-20


# Evaluation configurations
class EvaluationConfig:
    confusion_matrix = True
    per_class_metrics = True
    save_predictions = True


# Create a unique run ID for this training session
RUN_ID = datetime.now().strftime("%Y%m%d_%H%M%S")